
import cv2
import numpy as np
import matplotlib.pyplot as plt
import os


def calibrate_SIFT(img1, img2, show=True):
    img1 = np.array(img1).astype(np.float)
    img2 = np.array(img2).astype(np.float)

    img1 = np.uint8(img1 / np.mean(img1) * 50)
    img2 = np.uint8(img2 / np.mean(img2) * 50)


    '''
    img1 = cv2.Canny(img1, 180, 300).astype(np.uint8)
    img2 = cv2.Canny(img2, 180, 300).astype(np.uint8)
    '''
    # sift = cv2.xfeatures2d.SIFT_create()
    sift = cv2.SIFT_create()
    kp1, des1 = sift.detectAndCompute(img1, None)
    kp2, des2 = sift.detectAndCompute(img2, None)

    if show:
        plt.imshow(img1)
        plt.show()
        plt.imshow(img2)
        plt.show()
        plt.imshow(cv2.drawKeypoints(img2, kp2, None, color=(0, 255, 0), flags=0))
        plt.show()
        plt.imshow(cv2.drawKeypoints(img1, kp1, None, color=(0, 255, 0), flags=0))
        plt.show()

    # create BFMatcher object
    bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
    # Match descriptors.
    matches = bf.match(des1, des2)
    # Sort them in the order of their distance.
    matches = sorted(matches, key=lambda x: x.distance)

    goodMatch = matches[:20]
    ptsA = np.float32([kp1[m.queryIdx].pt for m in goodMatch]).reshape(-1, 1, 2)
    ptsB = np.float32([kp2[m.trainIdx].pt for m in goodMatch]).reshape(-1, 1, 2)
    # ptsA, ptsB = findPointsIn4Dir(matches, kp1, kp2, img1.shape)

    ransacReprojThreshold = 4
    H, status = cv2.estimateAffinePartial2D(ptsA, ptsB, cv2.RANSAC, ransacReprojThreshold)

    imgOut = cv2.warpAffine(img2, H, (img1.shape[1], img1.shape[0]), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP)
    diff = np.abs(cv2.Laplacian(img1, cv2.CV_64F) - cv2.Laplacian(imgOut, cv2.CV_64F))
    # diff = np.abs(img1 - imgOut)

    if show:
        # 显示对比
        plt.subplot(121)
        plt.title('SIFT')
        plt.imshow(img1)

        plt.subplot(122)
        plt.title('imgOut')
        plt.imshow(imgOut)

        plt.show()

    return H, diff


def findPointsIn4Dir(matches, pts1, pts2, imgsize, maxsize=20):
    res1 = [[], [], [], []]
    res2 = [[], [], [], []]
    for m in matches:
        if m.queryIdx < len(pts1) and m.queryIdx < len(pts2):
            kpt = pts1[m.queryIdx]
            if len(res1[0]) < maxsize / 4 and kpt.pt[0] < imgsize[0] / 2 and kpt.pt[1] < imgsize[1] / 2:
                res1[0].append(list(pts1[m.queryIdx].pt))
                res2[0].append(list(pts2[m.trainIdx].pt))
            elif len(res1[1]) < maxsize / 4 and kpt.pt[0] < imgsize[0] / 2 and kpt.pt[1] > imgsize[1] / 2:
                res1[1].append(list(pts1[m.queryIdx].pt))
                res2[1].append(list(pts2[m.trainIdx].pt))
            elif len(res1[2]) < maxsize / 4 and kpt.pt[0] > imgsize[0] / 2 and kpt.pt[1] < imgsize[1] / 2:
                res1[2].append(list(pts1[m.queryIdx].pt))
                res2[2].append(list(pts2[m.trainIdx].pt))
            elif len(res1[3]) < maxsize / 4 and kpt.pt[0] > imgsize[0] / 2 and kpt.pt[1] > imgsize[1] / 2:
                res1[3].append(list(pts1[m.queryIdx].pt))
                res2[3].append(list(pts2[m.trainIdx].pt))

    print(len(res1[0]), len(res1[1]), len(res1[2]), len(res1[3]))

    resA = []
    resB = []
    for i in range(len(res1)):
        resA.extend(res1[i])
        resB.extend(res2[i])

    return np.float32(resA), np.float32(resB)


def warpAffine(img, H):
    imgOut = cv2.warpAffine(img, H, (img.shape[1], img.shape[0]), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP)
    return imgOut


def readImgAs(path, resolution=(1000, 500), flip=False):
    img = cv2.imread(path)
    img1 = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    if flip:
        img2 = np.flip(img1, axis=1)
    else:
        img2 = img1
    return cv2.resize(img2, resolution)


def calAverageCalibration(imglist1, imglist2, show=True):
    H_mats = []
    for i in range(len(imglist1)):
        img_path1 = imglist1[i]
        img_path2 = imglist2[i]

        img1 = readImgAs(img_path1, flip=False)
        img2 = readImgAs(img_path2, flip=True)

        H_mat, diff = calibrate_SIFT(img1, img2, show=show)
        if show:
            plt.subplot(221)
            plt.imshow(img1 * 3)
            plt.subplot(222)
            plt.imshow(img2 * 3)
            plt.subplot(223)
            plt.imshow(warpAffine(img2, H_mat) * 3)
            plt.subplot(224)
            plt.imshow(diff, cmap='gray')

            plt.show()

        # print(np.mean(diff))
        if np.mean(diff) < 30:
            H_mats.append(H_mat)

    return np.mean(np.array(H_mats), axis=0), H_mats


if __name__ == "__main__":

    SHOWIMG = False

    # scene_root_dir = "G:/HIK_IMG_DIR/Indoors"
    scene_root_dir = "D:/HIK_IMG_DIR"

    imgdir_path_cam0 = scene_root_dir + '/USB[0]00F58881811/2021_04_25_17_56_08'
    imgdir_path_cam1 = scene_root_dir + '/USB[1]00F62123731/2021_04_25_17_56_08'

    imgs_cam0 = [imgdir_path_cam0 + "/" + filename for filename in os.listdir(imgdir_path_cam0)]
    imgs_cam1 = [imgdir_path_cam1 + "/" + filename for filename in os.listdir(imgdir_path_cam1)]

    H_mat = np.mat([[0.99479912, -0.01872361, -14.51830288],
                    [0.01872361, 0.99479912, -10.21850662]])

    H_mat, H_mats = calAverageCalibration(imgs_cam0[:10], imgs_cam1[:10], show=SHOWIMG)
    # print(H_mats)
    print(H_mat)

    np.save(scene_root_dir + "/cal_mat.npy", H_mat)

    img_path1 = imgs_cam0[12]
    img_path2 = imgs_cam1[12]

    img1 = readImgAs(img_path1)
    img2 = np.flip(readImgAs(img_path2), axis=1)

    img2_cali = warpAffine(img2, H_mat)

    cv2.imwrite("12_ref.png", img1 / np.mean(img1) * 50)
    cv2.imwrite("11_ori.png", img2 / np.mean(img2) * 50)
    cv2.imwrite("13_cali.png", img2_cali / np.mean(img2_cali) * 50)
